// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <memory>
#include <utility>

#include <turbo/utility/status.h>
#include <nebula/future/cancel.h>
#include <nebula/future/future.h>
#include <turbo/functional/functional.h>
#include <turbo/base/macros.h>
#include <nebula/util/type_fwd.h>
#include <turbo/base/macros.h>

namespace nebula::internal {

    /// \brief A group of related tasks
    ///
    /// A TaskGroup executes tasks with the signature `turbo::Status()`.
    /// Execution can be serial or parallel, depending on the TaskGroup
    /// implementation.  When finish() returns, it is guaranteed that all
    /// tasks have finished, or at least one has errored.
    ///
    /// Once an error has occurred any tasks that are submitted to the task group
    /// will not run.  The call to append will simply return without scheduling the
    /// task.
    ///
    /// If the task group is parallel it is possible that multiple tasks could be
    /// running at the same time and one of those tasks fails.  This will put the
    /// task group in a failure state (so additional tasks cannot be run) however
    /// it will not interrupt running tasks.  finish will not complete
    /// until all running tasks have finished, even if one task fails.
    ///
    /// Once a task group has finished new tasks may not be added to it.  If you need to start
    /// a new batch of work then you should create a new task group.
    class TURBO_EXPORT TaskGroup : public std::enable_shared_from_this<TaskGroup> {
    public:
        /// Add a turbo::Status-returning function to execute.  Execution order is
        /// undefined.  The function may be executed immediately or later.
        template<typename Function>
        void append(Function &&func) {
            return AppendReal(std::forward<Function>(func));
        }

        /// Wait for execution of all tasks (and subgroups) to be finished,
        /// or for at least one task (or subgroup) to error out.
        /// The returned turbo::Status propagates the error status of the first failing
        /// task (or subgroup).
        virtual turbo::Status finish() = 0;

        /// Returns a future that will complete the first time all tasks are finished.
        /// This should be called only after all top level tasks
        /// have been added to the task group.
        ///
        /// If you are using a TaskGroup asynchronously there are a few considerations to keep
        /// in mind.  The tasks should not block on I/O, etc (defeats the purpose of using
        /// futures) and should not be doing any nested locking or you run the risk of the tasks
        /// getting stuck in the thread pool waiting for tasks which cannot get scheduled.
        ///
        /// Primarily this call is intended to help migrate existing work written with TaskGroup
        /// in mind to using futures without having to do a complete conversion on the first
        /// pass.
        virtual Future<> FinishAsync() = 0;

        /// The current aggregate error turbo::Status.  Non-blocking, useful for stopping early.
        virtual turbo::Status current_status() = 0;

        /// Whether some tasks have already failed.  Non-blocking, useful for stopping early.
        virtual bool ok() const = 0;

        /// How many tasks can typically be executed in parallel.
        /// This is only a hint, useful for testing or debugging.
        virtual int parallelism() = 0;

        static std::shared_ptr<TaskGroup> MakeSerial(StopToken = StopToken::unstoppable());

        static std::shared_ptr<TaskGroup> MakeThreaded(internal::Executor *,
                                                       StopToken = StopToken::unstoppable());

        virtual ~TaskGroup() = default;

    protected:
        TaskGroup() = default;
        TURBO_DISALLOW_COPY_AND_ASSIGN(TaskGroup);

        virtual void AppendReal(turbo::FnOnce<turbo::Status()> task) = 0;
    };

}  // namespace nebula::internal
